# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time        : 2020/11/23 22:44
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : Chapter2.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""

import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import solve
from scipy import interpolate
from scipy.interpolate import lagrange

plt.rc('figure', figsize=(20, 15))

print("Problem I:")

given_x = [0.2, 0.4, 0.6, 0.8, 1.0]
given_y = [0.98, 0.92, 0.81, 0.64, 0.38]
given_times = 4
x_range = (0, 1.1, 0.02)


# @brief: Convert(begin,end,interval) to a list, but interval can be float numbers.
def process_xpara(xpara):
    max_times = 0
    if 0 < xpara[2] < 1:
        tmp_xpara_interval = xpara[2]
        while tmp_xpara_interval - int(tmp_xpara_interval) != 0:
            max_times = max_times + 1
            tmp_xpara_interval = tmp_xpara_interval * 10
    max_times = 10 ** max_times
    return [i / max_times for i in
            range(int(xpara[0] * max_times), int(xpara[1] * max_times), int(xpara[2] * max_times))]


def divide_difference(x, y, times):
    now = [(x[i], y[i]) for i in range(len(x))]
    ans = [now[0]]
    for order in range(1, times + 1):
        tmp = []
        for i in range(1, len(now)):
            tmp.append((x[order + i - 1] - x[i - 1], (now[i][1] - now[i - 1][1]) / (x[order + i - 1] - x[i - 1])))
        now = tmp
        ans.append(now[0])
    return ans


def get_func_value_newton(xcoef, x, xorigin):
    ans = 0
    for i in range(len(xcoef)):
        tmp = xcoef[i][1]
        for j in range(i):
            tmp = tmp * (x - xorigin[j])
        ans = ans + tmp
    return ans


"""
#@param: xpara(xbegin,xend,xinterval) fpara[f[x_1~i]]
"""


# spec_i=[0.2+0.08*x for x in (0,1,10,11)]

def newton_interpolate(xpara, fpara, xorigin):
    x_discrete_value = process_xpara(xpara)
    return [(x, get_func_value_newton(fpara, x, xorigin)) for x in x_discrete_value]


parameters = divide_difference(given_x, given_y, given_times)
newton_interpolate_value = newton_interpolate(x_range, parameters, given_x)

fig = plt.figure()
sub_fig1 = fig.add_subplot(2, 2, 1)
sub_fig1.set_title("Problem I")
sub_fig1.plot([var[0] for var in newton_interpolate_value], [var[1] for var in newton_interpolate_value],
              label='Newton')

# l_f=lagrange(given_x,given_y)
# tmpara=process_xpara(x_range)
# plt.plot(tmpara,[l_f(x) for x in tmpara])

# 三次样条插值
n = len(given_x)
h = []
f0p = 0
fnp = 0
for i in range(1, len(given_x)):
    h.append(given_x[i] - given_x[i - 1])
miu = [0]  # 0 should not be used
lam = [1]
d = [6 / h[0] * ((given_y[1] - given_y[0]) / (given_x[1] - given_x[0]) - f0p)]
for j in range(1, len(h)):
    miu.append(h[j - 1] / (h[j - 1] + h[j]))
    lam.append(h[j] / (h[j - 1] + h[j]))
    d.append(6 * ((given_y[j + 1] - given_y[j]) / (given_x[j + 1] - given_x[j]) - (given_y[j - 1] - given_y[j]) / (
                given_x[j - 1] - given_x[j])) / (h[j - 1] + h[j]))
miu.append(1)
d.append(6 / h[-1] * (fnp - (given_y[-1] - given_y[-2]) / (given_x[-1] - given_x[-2])))

A = np.zeros((n, n))
for i in range(n):
    A[i][i] = 2
    if i != n - 1:
        A[i][i + 1] = lam[i]
    if i != 0:
        A[i][i - 1] = miu[i]
C = solve(A, np.array(d).T)


# print(C)

def get_func_value_cubic_spline(mtuple, xtuple, ytuple, x):
    return mtuple[0] / (6 * (xtuple[1] - xtuple[0])) * (xtuple[1] - x) ** 3 + mtuple[1] / (
                6 * (xtuple[1] - xtuple[0])) * (x - xtuple[0]) ** 3 + (
                       ytuple[0] - (mtuple[0] * (xtuple[1] - xtuple[0]) ** 2 / 6)) * (xtuple[1] - x) / (
                       xtuple[1] - xtuple[0]) + (ytuple[1] - (mtuple[1] * (xtuple[1] - xtuple[0]) ** 2 / 6)) * (
                       x - xtuple[0]) / (xtuple[1] - xtuple[0])


def cubic_spline_interpolate(xpara, mpara, x, y):
    fun_value = []
    x_discrete_value = process_xpara(xpara)
    for j in range(len(x) - 1):
        ok_value = [(element,
                     get_func_value_cubic_spline((mpara[j], mpara[j + 1]), (x[j], x[j + 1]), (y[j], y[j + 1]), element))
                    for element in x_discrete_value if x[j] <= element < x[j + 1]]
        fun_value = fun_value + ok_value
    return fun_value


cubic_spline_interpolate_value = cubic_spline_interpolate(x_range, C.tolist(), given_x, given_y)

sub_fig1.plot([var[0] for var in cubic_spline_interpolate_value], [var[1] for var in cubic_spline_interpolate_value],
              label='Cubic')

sub_fig1.legend(loc='best')


def get_func_x(x):
    return 1 / (1 + 25 * x * x)


given_x = np.linspace(-1, 1, 10)
given_y = get_func_x(given_x)  # p.array([get_func_x(x) for x in given_x])
display_x = np.linspace(-1, 1, 100)
display_y = get_func_x(display_x)

sub_fig2 = fig.add_subplot(2, 2, 2)
sub_fig2.set_title("Problem II(Alpha): Using System Functions")

c_x = interpolate.interp1d(given_x, given_y, kind='cubic')
l_x = lagrange(given_x, given_y)
sub_fig2.plot(display_x, l_x(display_x))
sub_fig2.plot(display_x, c_x(display_x))
sub_fig2.plot(display_x, display_y)

sub_fig3 = fig.add_subplot(2, 2, 3)
sub_fig3.set_title("Problem II(Beta): Using System Functions")

given_x = np.linspace(-1, 1, 20)
given_y = get_func_x(given_x)  # p.array([get_func_x(x) for x in given_x])
c_x = interpolate.interp1d(given_x, given_y, kind='cubic')
l_x = lagrange(given_x, given_y)
sub_fig3.plot(display_x, l_x(display_x))
sub_fig3.plot(display_x, c_x(display_x))
sub_fig3.plot(display_x, display_y)

fig_problem_three = plt.figure()

given_x = [0, 1, 4, 9, 16, 25, 36, 49, 64]
given_y = [0, 1, 2, 3, 4, 5, 6, 7, 8]
display_big_x = np.linspace(0, 64, 200)
display_small_x = np.linspace(0, 1, 50)
sub_fig4 = fig_problem_three.add_subplot(2, 1, 1)
l_x = lagrange(given_x, given_y)
c_x = interpolate.interp1d(given_x, given_y, kind='cubic')
sub_fig4.plot(display_big_x, l_x(display_big_x), label='Lagrange')
sub_fig4.plot(display_big_x, c_x(display_big_x), label='Cubic')
sub_fig4.plot(display_big_x, np.sqrt(display_big_x), label='Origin')
sub_fig4.legend(loc='best')

sub_fig5 = fig_problem_three.add_subplot(2, 1, 2)
sub_fig5.plot(display_small_x, l_x(display_small_x), label='Lagrange')
sub_fig5.plot(display_small_x, c_x(display_small_x), label='Cubic')
sub_fig5.plot(display_small_x, np.sqrt(display_small_x), label='Origin')
sub_fig5.legend(loc='best')
